{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ccfefb2a",
   "metadata": {},
   "source": [
    "给高loss的数据加权,提高这些数据被采样中的概率.\n",
    "\n",
    "为了缓解过拟合,降低这些数据的lr,这里以削减loss的方式实现."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91162df8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/P9b71AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAS2ElEQVR4nO3df0xT578H8DcFWlA4Zehoh9BovjNT4q8NFY7mZot2Mke8c5KbbXGOGaPRFa/KrtlI/DHdFoz7w/1S/GdT/3EuLHGLXH+MoWI2q2gdCYKSmetSvmpbf1xOgY1S6HP/+IZz1wnOYu1D3fuVnMTzPJ+2n/Ng37TnUEgQQggQEUlgkN0AEf19MYCISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaaQG0c+dOjB07FikpKSgoKEBDQ4OsVohIEikB9PXXX6O8vBybN2/GhQsXMHXqVBQVFcHn88loh4gkSZDxYdSCggLMmDEDn3/+OQAgFAohNzcXq1evxrvvvhvrdohIkqRYP2BPTw9cLhcqKir0MYPBALvdDqfTOeBtAoEAAoGAvh8KhXDnzh2MGjUKCQkJD71nIoqMEAIdHR3Izs6GwTD4G62YB9CtW7fQ19cHi8USNm6xWHD58uUBb1NZWYktW7bEoj0iiqK2tjbk5OQMOh/zABqKiooKlJeX6/uapsFms6GtrQ2KokjsjIgG4vf7kZubi/T09HvWxTyARo8ejcTERHi93rBxr9cLq9U64G1MJhNMJtNd44qiMICIhrG/OkUS86tgRqMR+fn5qKur08dCoRDq6uqgqmqs2yEiiaS8BSsvL0dpaSmmT5+OmTNn4uOPP0ZXVxeWLl0qox0ikkRKAL3yyiu4efMmNm3aBI/Hg2nTpuHo0aN3nZgmokeblJ8DelB+vx9msxmapvEcENEwdL/PUX4WjIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNJEHECnTp3CggULkJ2djYSEBHz77bdh80IIbNq0CU888QRSU1Nht9vxyy+/hNXcuXMHixcvhqIoyMjIwLJly9DZ2flAB0JE8SfiAOrq6sLUqVOxc+fOAee3b9+OTz/9FLt378bZs2cxcuRIFBUVobu7W69ZvHgxmpubUVtbi5qaGpw6dQorVqwY+lEQUXwSDwCAOHjwoL4fCoWE1WoVH330kT7W3t4uTCaT+Oqrr4QQQrS0tAgA4ty5c3rNkSNHREJCgrh27dp9Pa6maQKA0DTtQdonoofkfp+jUT0HdPXqVXg8Htjtdn3MbDajoKAATqcTAOB0OpGRkYHp06frNXa7HQaDAWfPnh3wfgOBAPx+f9hGRPEvqgHk8XgAABaLJWzcYrHocx6PB1lZWWHzSUlJyMzM1Gv+rLKyEmazWd9yc3Oj2TYRSRIXV8EqKiqgaZq+tbW1yW6JiKIgqgFktVoBAF6vN2zc6/Xqc1arFT6fL2y+t7cXd+7c0Wv+zGQyQVGUsI2I4l9UA2jcuHGwWq2oq6vTx/x+P86ePQtVVQEAqqqivb0dLpdLrzl+/DhCoRAKCgqi2Q4RDXNJkd6gs7MTV65c0fevXr2KxsZGZGZmwmazYe3atfjggw8wfvx4jBs3Dhs3bkR2djYWLlwIAJg4cSJeeOEFLF++HLt370YwGERZWRleffVVZGdnR+3AiCgORHp57cSJEwLAXVtpaakQ4l+X4jdu3CgsFoswmUxi7ty5orW1New+bt++LV577TWRlpYmFEURS5cuFR0dHVG/xEdEctzvczRBCCEk5t+Q+P1+mM1maJrG80FEw9D9Pkfj4ioYET2aGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNBH/WR6ih+G32//E73eu6ftJqelQxkxEQkKCxK7oYWMA0bBw539cuHHhv/X9NOt4KGMmSuyIYoFvwYhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0EQVQZWUlZsyYgfT0dGRlZWHhwoVobW0Nq+nu7obD4cCoUaOQlpaGkpISeL3esBq3243i4mKMGDECWVlZWL9+PXp7ex/8aIgorkQUQPX19XA4HDhz5gxqa2sRDAYxb948dHV16TXr1q3DoUOHUF1djfr6ely/fh2LFi3S5/v6+lBcXIyenh6cPn0a+/btw969e7Fp06boHRURxQfxAHw+nwAg6uvrhRBCtLe3i+TkZFFdXa3XXLp0SQAQTqdTCCHE4cOHhcFgEB6PR6+pqqoSiqKIQCBwX4+raZoAIDRNe5D2aRhpa/hWNOxerm8t324XoVBIdls0RPf7HH2gc0CapgEAMjMzAQAulwvBYBB2u12vmTBhAmw2G5xOJwDA6XRi8uTJsFgsek1RURH8fj+am5sHfJxAIAC/3x+2EVH8G3IAhUIhrF27FrNnz8akSZMAAB6PB0ajERkZGWG1FosFHo9Hr/lj+PTP988NpLKyEmazWd9yc3OH2jYRDSNDDiCHw4GLFy/iwIED0exnQBUVFdA0Td/a2toe+mMS0cM3pA+jlpWVoaamBqdOnUJOTo4+brVa0dPTg/b29rBXQV6vF1arVa9paGgIu7/+q2T9NX9mMplgMpmG0ioRDWMRvQISQqCsrAwHDx7E8ePHMW7cuLD5/Px8JCcno66uTh9rbW2F2+2GqqoAAFVV0dTUBJ/Pp9fU1tZCURTk5eU9yLEQUZyJ6BWQw+HA/v378d133yE9PV0/Z2M2m5Gamgqz2Yxly5ahvLwcmZmZUBQFq1evhqqqKCwsBADMmzcPeXl5WLJkCbZv3w6Px4MNGzbA4XDwVQ7R30xEAVRVVQUAeO6558LG9+zZgzfffBMAsGPHDhgMBpSUlCAQCKCoqAi7du3SaxMTE1FTU4NVq1ZBVVWMHDkSpaWl2Lp164MdCRHFnQQhhJDdRKT8fj/MZjM0TYOiKLLboSj457nv7vqFZBP+/b/4GxHj1P0+R/lZMCKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQR/WVUoqEKBoPo6uoadD7Q3R2239fXC629HRjkDxOaTCakpqZGs0WSgAFEMfHTTz/h9ddfH3T+tWfH4z/+bby+f/68C/P/czIG+7O9DocDFRUVUe6SYo0BRDERCARw7dq1Qee1jifQFpiAG4F/wJx0E12BI7h27dqgAeT3+x9OoxRTEZ0DqqqqwpQpU6AoChRFgaqqOHLkiD7f3d0Nh8OBUaNGIS0tDSUlJfB6vWH34Xa7UVxcjBEjRiArKwvr169Hb29vdI6G4taNwJO41DkLd4JjcPX3qWjtKhg0fOjREVEA5eTkYNu2bXC5XDh//jzmzJmDl156Cc3NzQCAdevW4dChQ6iurkZ9fT2uX7+ORYsW6bfv6+tDcXExenp6cPr0aezbtw979+7Fpk2bontUFHe6+swI6S/IE9DZlwFg4PM/9OiI6C3YggULwvY//PBDVFVV4cyZM8jJycEXX3yB/fv3Y86cOQCAPXv2YOLEiThz5gwKCwvx/fffo6WlBT/88AMsFgumTZuG999/H++88w7ee+89GI3G6B0ZxZUsoxsphg50h9KQmBDEGNMVgK+BHnlDPgfU19eH6upqdHV1QVVVuFwuBINB2O12vWbChAmw2WxwOp0oLCyE0+nE5MmTYbFY9JqioiKsWrUKzc3NePrppyPq4fLly0hLSxvqIVAMud3ue863/vIzgsGtuN0zBmlJ/4vezl/uWX/r1i20tLREs0WKos7OzvuqiziAmpqaoKoquru7kZaWhoMHDyIvLw+NjY0wGo3IyMgIq7dYLPB4PAAAj8cTFj798/1zgwkEAggEAvp+/wlITdN4/ihO3OsSPAA0XvGg8crg/wf+LBAIoL29/QG7ooflr77e/SIOoKeeegqNjY3QNA3ffPMNSktLUV9fH3GDkaisrMSWLVvuGi8oKICiKA/1sSk6Ojo6onp/Y8aMwaxZs6J6nxQ993uVMuKfhDYajXjyySeRn5+PyspKTJ06FZ988gmsVit6enru+q7k9XphtVoBAFar9a6rYv37/TUDqaiogKZp+tbW1hZp20Q0DD3wRzFCoRACgQDy8/ORnJyMuro6fa61tRVutxuqqgIAVFVFU1MTfD6fXlNbWwtFUZCXlzfoY5hMJv3Sf/9GRPEvordgFRUVmD9/Pmw2Gzo6OrB//36cPHkSx44dg9lsxrJly1BeXo7MzEwoioLVq1dDVVUUFhYCAObNm4e8vDwsWbIE27dvh8fjwYYNG+BwOGAymR7KARLR8BVRAPl8Przxxhu4ceMGzGYzpkyZgmPHjuH5558HAOzYsQMGgwElJSUIBAIoKirCrl279NsnJiaipqYGq1atgqqqGDlyJEpLS7F169boHhUNO0lJSVF95cpvWI+GBCFE3P2whd/vh9lshqZpfDsWJ7q7u3Hz5s2o3V96evpdV1xp+Ljf5yg/C0YxkZKSgtzcXNlt0DDD3wdERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpEmS3cBQCCEAAH6/X3InRDSQ/udm/3N1MHEZQLdv3wYA5ObmSu6EiO6lo6MDZrN50Pm4DKDMzEwAgNvtvufBUTi/34/c3Fy0tbVBURTZ7cQFrtnQCCHQ0dGB7Ozse9bFZQAZDP86dWU2m/mfYggUReG6RYhrFrn7eXHAk9BEJA0DiIikicsAMplM2Lx5M0wmk+xW4grXLXJcs4crQfzVdTIioockLl8BEdGjgQFERNIwgIhIGgYQEUkTlwG0c+dOjB07FikpKSgoKEBDQ4PslqSprKzEjBkzkJ6ejqysLCxcuBCtra1hNd3d3XA4HBg1ahTS0tJQUlICr9cbVuN2u1FcXIwRI0YgKysL69evR29vbywPRZpt27YhISEBa9eu1ce4ZjEi4syBAweE0WgUX375pWhubhbLly8XGRkZwuv1ym5NiqKiIrFnzx5x8eJF0djYKF588UVhs9lEZ2enXrNy5UqRm5sr6urqxPnz50VhYaGYNWuWPt/b2ysmTZok7Ha7+Pnnn8Xhw4fF6NGjRUVFhYxDiqmGhgYxduxYMWXKFLFmzRp9nGsWG3EXQDNnzhQOh0Pf7+vrE9nZ2aKyslJiV8OHz+cTAER9fb0QQoj29naRnJwsqqur9ZpLly4JAMLpdAohhDh8+LAwGAzC4/HoNVVVVUJRFBEIBGJ7ADHU0dEhxo8fL2pra8Wzzz6rBxDXLHbi6i1YT08PXC4X7Ha7PmYwGGC32+F0OiV2Nnxomgbg/z+w63K5EAwGw9ZswoQJsNls+po5nU5MnjwZFotFrykqKoLf70dzc3MMu48th8OB4uLisLUBuGaxFFcfRr116xb6+vrCvugAYLFYcPnyZUldDR+hUAhr167F7NmzMWnSJACAx+OB0WhERkZGWK3FYoHH49FrBlrT/rlH0YEDB3DhwgWcO3furjmuWezEVQDRvTkcDly8eBE//vij7FaGtba2NqxZswa1tbVISUmR3c7fWly9BRs9ejQSExPvuhrh9XphtVoldTU8lJWVoaamBidOnEBOTo4+brVa0dPTg/b29rD6P66Z1WodcE375x41LpcLPp8PzzzzDJKSkpCUlIT6+np8+umnSEpKgsVi4ZrFSFwFkNFoRH5+Purq6vSxUCiEuro6qKoqsTN5hBAoKyvDwYMHcfz4cYwbNy5sPj8/H8nJyWFr1traCrfbra+ZqqpoamqCz+fTa2pra6EoCvLy8mJzIDE0d+5cNDU1obGxUd+mT5+OxYsX6//mmsWI7LPgkTpw4IAwmUxi7969oqWlRaxYsUJkZGSEXY34O1m1apUwm83i5MmT4saNG/r222+/6TUrV64UNptNHD9+XJw/f16oqipUVdXn+y8pz5s3TzQ2NoqjR4+Kxx9//G91SfmPV8GE4JrFStwFkBBCfPbZZ8Jmswmj0Shmzpwpzpw5I7slaQAMuO3Zs0ev+f3338Vbb70lHnvsMTFixAjx8ssvixs3boTdz6+//irmz58vUlNTxejRo8Xbb78tgsFgjI9Gnj8HENcsNvjrOIhImrg6B0REjxYGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJM3/AQFSMa4NZCopAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "\n",
    "#定义环境\n",
    "class MyWrapper(gym.Wrapper):\n",
    "\n",
    "    def __init__(self):\n",
    "        env = gym.make('CartPole-v1', render_mode='rgb_array')\n",
    "        super().__init__(env)\n",
    "        self.env = env\n",
    "        self.step_n = 0\n",
    "\n",
    "    def reset(self):\n",
    "        state, _ = self.env.reset()\n",
    "        self.step_n = 0\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        state, reward, terminated, truncated, info = self.env.step(action)\n",
    "        over = terminated or truncated\n",
    "\n",
    "        #限制最大步数\n",
    "        self.step_n += 1\n",
    "        if self.step_n >= 200:\n",
    "            over = True\n",
    "        \n",
    "        #没坚持到最后,扣分\n",
    "        if over and self.step_n < 200:\n",
    "            reward = -1000\n",
    "\n",
    "        return state, reward, over\n",
    "\n",
    "    #打印游戏图像\n",
    "    def show(self):\n",
    "        from matplotlib import pyplot as plt\n",
    "        plt.figure(figsize=(3, 3))\n",
    "        plt.imshow(self.env.render())\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "env = MyWrapper()\n",
    "\n",
    "env.reset()\n",
    "\n",
    "env.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ecfbe912",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Sequential(\n",
       "   (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "   (1): ReLU()\n",
       "   (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "   (3): ReLU()\n",
       "   (4): Linear(in_features=64, out_features=2, bias=True)\n",
       " ),\n",
       " Sequential(\n",
       "   (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "   (1): ReLU()\n",
       "   (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "   (3): ReLU()\n",
       "   (4): Linear(in_features=64, out_features=2, bias=True)\n",
       " ))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "#定义模型,评估状态下每个动作的价值\n",
    "model = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    ")\n",
    "\n",
    "#延迟更新的模型,用于计算target\n",
    "model_delay = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    ")\n",
    "\n",
    "#复制参数\n",
    "model_delay.load_state_dict(model.state_dict())\n",
    "\n",
    "model, model_delay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84cbf0ff",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-973.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython import display\n",
    "import random\n",
    "\n",
    "\n",
    "#玩一局游戏并记录数据\n",
    "def play(show=False):\n",
    "    data = []\n",
    "    reward_sum = 0\n",
    "\n",
    "    state = env.reset()\n",
    "    over = False\n",
    "    while not over:\n",
    "        action = model(torch.FloatTensor(state).reshape(1, 4)).argmax().item()\n",
    "        if random.random() < 0.1:\n",
    "            action = env.action_space.sample()\n",
    "\n",
    "        next_state, reward, over = env.step(action)\n",
    "\n",
    "        data.append((state, action, reward, next_state, over))\n",
    "        reward_sum += reward\n",
    "\n",
    "        state = next_state\n",
    "\n",
    "        if show:\n",
    "            display.clear_output(wait=True)\n",
    "            env.show()\n",
    "\n",
    "    return data, reward_sum\n",
    "\n",
    "\n",
    "play()[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c79f77f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_7051/919455838.py:35: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at  ../torch/csrc/utils/tensor_new.cpp:201.)\n",
      "  state = torch.FloatTensor([i[0] for i in data]).reshape(-1, 4)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(201,\n",
       " (array([0.0444518 , 0.03310756, 0.03021159, 0.02350001], dtype=float32),\n",
       "  1,\n",
       "  1.0,\n",
       "  array([ 0.04511395,  0.22778352,  0.03068159, -0.25949982], dtype=float32),\n",
       "  False))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#数据池\n",
    "class Pool:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.pool = []\n",
    "        self.prob = []\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.pool)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self.pool[i]\n",
    "\n",
    "    #更新动作池\n",
    "    def update(self):\n",
    "        #每次更新不少于N条新数据\n",
    "        old_len = len(self.pool)\n",
    "        while len(pool) - old_len < 200:\n",
    "            data = play()[0]\n",
    "            self.pool.extend(data)\n",
    "            #维护概率表\n",
    "            self.prob.extend([1.0] * len(data))\n",
    "\n",
    "        #只保留最新的N条数据\n",
    "        self.pool = self.pool[-2_0000:]\n",
    "        self.prob = self.prob[-2_0000:]\n",
    "\n",
    "    #获取一批数据样本\n",
    "    def sample(self):\n",
    "        idx = torch.FloatTensor(self.prob).clamp(0.1, 1.0).multinomial(\n",
    "            num_samples=64, replacement=False)\n",
    "\n",
    "        data = [self.pool[i] for i in idx]\n",
    "\n",
    "        state = torch.FloatTensor([i[0] for i in data]).reshape(-1, 4)\n",
    "        action = torch.LongTensor([i[1] for i in data]).reshape(-1, 1)\n",
    "        reward = torch.FloatTensor([i[2] for i in data]).reshape(-1, 1)\n",
    "        next_state = torch.FloatTensor([i[3] for i in data]).reshape(-1, 4)\n",
    "        over = torch.LongTensor([i[4] for i in data]).reshape(-1, 1)\n",
    "\n",
    "        return idx, state, action, reward, next_state, over\n",
    "\n",
    "\n",
    "pool = Pool()\n",
    "pool.update()\n",
    "pool.sample()\n",
    "\n",
    "len(pool), pool[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "018c4b1e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 412 [0.8312286138534546] -990.4\n",
      "100 20000 [0.7296427488327026, 1.0, 0.7277626991271973, 1.0] -416.75\n",
      "200 20000 [0.9995549321174622, 0.9999997615814209, 0.9968600273132324, 1.0] -105.65\n",
      "300 20000 [0.5885254144668579, 0.5119339227676392, 0.9999974966049194, 1.0] 200.0\n",
      "400 20000 [0.5389884114265442, 0.7463151812553406, 0.6541365385055542, 0.9985623955726624] 200.0\n",
      "500 20000 [0.5640982389450073, 0.530198335647583, 0.5991714000701904, 0.5062503814697266] 200.0\n",
      "600 20000 [0.7470700740814209, 0.5565980076789856, 0.9851357340812683, 0.5000351667404175] 200.0\n",
      "700 20000 [0.5007219910621643, 0.5155380964279175, 0.5008571147918701, 0.5083107352256775] 143.0\n",
      "800 20000 [0.5586060881614685, 0.520609974861145, 0.5244655013084412, 1.0] 200.0\n",
      "900 20000 [0.5362679362297058, 0.5268680453300476, 0.5999729633331299, 0.5125797390937805] 200.0\n"
     ]
    }
   ],
   "source": [
    "#训练\n",
    "def train():\n",
    "    model.train()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)\n",
    "    loss_fn = torch.nn.MSELoss(reduction='none')\n",
    "\n",
    "    #共更新N轮数据\n",
    "    for epoch in range(1000):\n",
    "        pool.update()\n",
    "\n",
    "        #每次更新数据后,训练N次\n",
    "        for i in range(200):\n",
    "\n",
    "            #采样N条数据\n",
    "            idx, state, action, reward, next_state, over = pool.sample()\n",
    "\n",
    "            #计算value\n",
    "            value = model(state).gather(dim=1, index=action)\n",
    "\n",
    "            #计算target\n",
    "            with torch.no_grad():\n",
    "                target = model_delay(next_state)\n",
    "            target = target.max(dim=1)[0].reshape(-1, 1)\n",
    "            target = target * 0.99 * (1 - over) + reward\n",
    "\n",
    "            #根据概率缩放loss\n",
    "            r = torch.FloatTensor([pool.prob[i] for i in idx])\n",
    "            r = (1 - r).clamp(0.1, 1.0).reshape(-1, 1)\n",
    "\n",
    "            loss = loss_fn(value, target)\n",
    "            (loss * r).mean(0).backward()\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            #根据loss调整数据权重\n",
    "            for i, j in zip(idx.tolist(),\n",
    "                            loss.abs().sigmoid().flatten().tolist()):\n",
    "                pool.prob[i] = j\n",
    "\n",
    "        #复制参数\n",
    "        if (epoch + 1) % 5 == 0:\n",
    "            model_delay.load_state_dict(model.state_dict())\n",
    "\n",
    "        if epoch % 100 == 0:\n",
    "            test_result = sum([play()[-1] for _ in range(20)]) / 20\n",
    "            print(epoch, len(pool), pool.prob[::5000], test_result)\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a101d0e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/P9b71AAAACXBIWXMAAA9hAAAPYQGoP6dpAAATiklEQVR4nO3df2xT5b8H8He7rYWxnc6Ba11YAwkoLvxQB4yjf2i0MmExosu9aohOQiBiR4AZcl2ieBXNDCb+QPnxjwKJQbwzmcYFf+xuMuJdYTJdMobsq/eLbgHainx32m2s7dbn/mF2rpUN19H16anvV3ISznmenn7Ow/reOc9pV5MQQoCISAKz7AKI6O+LAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNJIC6A9e/Zgzpw5mDZtGkpLS9HW1iarFCKSREoAffTRR6iursaLL76I7777DkuWLEFZWRn8fr+McohIEpOMD6OWlpZi2bJlePfddwEA0WgURUVF2Lx5M5577rlkl0NEkmQm+wnD4TDa29tRU1OjbzObzXC5XPB4PGM+JhQKIRQK6evRaBSXL1/GzJkzYTKZprxmIoqPEALBYBCFhYUwm8e/0Ep6AF26dAkjIyOw2+0x2+12O86ePTvmY2pra/HSSy8lozwiSqDe3l7Mnj173PakB9Bk1NTUoLq6Wl/XNA1OpxO9vb1QFEViZUQ0lkAggKKiIuTm5l6zX9IDaNasWcjIyIDP54vZ7vP54HA4xnyM1WqF1Wq9aruiKAwgohT2V1MkSb8LZrFYUFJSgqamJn1bNBpFU1MTVFVNdjlEJJGUS7Dq6mpUVlZi6dKlWL58Od566y0MDAxg3bp1MsohIkmkBNCjjz6KX3/9FTt27IDX68Vtt92GL7744qqJaSJKb1LeB3S9AoEAbDYbNE3jHBBRCproa5SfBSMiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0cQfQ8ePH8eCDD6KwsBAmkwmffPJJTLsQAjt27MBNN92E6dOnw+Vy4ccff4zpc/nyZaxduxaKoiAvLw/r169Hf3//dR0IERlP3AE0MDCAJUuWYM+ePWO279q1C7t378b+/ftx8uRJzJgxA2VlZRgaGtL7rF27Fl1dXWhsbERDQwOOHz+OjRs3Tv4oiMiYxHUAIOrr6/X1aDQqHA6HeP311/VtfX19wmq1ig8//FAIIcSZM2cEAPHtt9/qfT7//HNhMpnE+fPnJ/S8mqYJAELTtOspn4imyERfowmdAzp37hy8Xi9cLpe+zWazobS0FB6PBwDg8XiQl5eHpUuX6n1cLhfMZjNOnjw55n5DoRACgUDMQkTGl9AA8nq9AAC73R6z3W63621erxcFBQUx7ZmZmcjPz9f7/FltbS1sNpu+FBUVJbJsIpLEEHfBampqoGmavvT29souiYgSIKEB5HA4AAA+ny9mu8/n09scDgf8fn9M+/DwMC5fvqz3+TOr1QpFUWIWIjK+hAbQ3Llz4XA40NTUpG8LBAI4efIkVFUFAKiqir6+PrS3t+t9mpubEY1GUVpamshyiCjFZcb7gP7+fvz000/6+rlz59DR0YH8/Hw4nU5s3boVr7zyCubPn4+5c+fihRdeQGFhIdasWQMAuPXWW/HAAw9gw4YN2L9/PyKRCKqqqvDYY4+hsLAwYQdGRAYQ7+21r7/+WgC4aqmsrBRC/H4r/oUXXhB2u11YrVZx3333ie7u7ph9/Pbbb+Lxxx8XOTk5QlEUsW7dOhEMBhN+i4+I5Jjoa9QkhBAS829SAoEAbDYbNE3jfBBRCproa9QQd8GIKD0xgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJp4v5aHqJkCg/8C8EL/9DXzZlZsDkXw5zBH910wP9FSmkDv/bgn83v6etZ0xUs/PebYc7IkVgVJQovwSjFGe5boygODCBKbcb72jqKAwOIUpoBvzeT4sAAotQmorIroCnEAKKUJjgHlNYYQJTaeAmW1uIKoNraWixbtgy5ubkoKCjAmjVr0N3dHdNnaGgIbrcbM2fORE5ODioqKuDz+WL69PT0oLy8HNnZ2SgoKMD27dsxPDx8/UdDaYdzQOktrgBqaWmB2+3GiRMn0NjYiEgkgpUrV2JgYEDvs23bNnz22Weoq6tDS0sLLly4gEceeURvHxkZQXl5OcLhMFpbW3Ho0CEcPHgQO3bsSNxRUfrgHFB6E9fB7/cLAKKlpUUIIURfX5/IysoSdXV1ep8ffvhBABAej0cIIcTRo0eF2WwWXq9X77Nv3z6hKIoIhUITel5N0wQAoWna9ZRPBvDr2f8Rbfs36Mv3h54VkStB2WXRX5joa/S65oA0TQMA5OfnAwDa29sRiUTgcrn0PgsWLIDT6YTH4wEAeDweLFq0CHa7Xe9TVlaGQCCArq6uMZ8nFAohEAjELPT3IHgJltYmHUDRaBRbt27FXXfdhYULFwIAvF4vLBYL8vLyYvra7XZ4vV69zx/DZ7R9tG0stbW1sNls+lJUVDTZssloGEBpbdIB5Ha7cfr0aRw5ciSR9YyppqYGmqbpS29v75Q/J6UG3oZPb5P6MGpVVRUaGhpw/PhxzJ49W9/ucDgQDofR19cXcxbk8/ngcDj0Pm1tbTH7G71LNtrnz6xWK6xW62RKJaPjJHRai+sMSAiBqqoq1NfXo7m5GXPnzo1pLykpQVZWFpqamvRt3d3d6OnpgaqqAABVVdHZ2Qm/36/3aWxshKIoKC4uvp5joXTES7C0FtcZkNvtxuHDh/Hpp58iNzdXn7Ox2WyYPn06bDYb1q9fj+rqauTn50NRFGzevBmqqmLFihUAgJUrV6K4uBhPPPEEdu3aBa/Xi+effx5ut5tnOXQVTkKnt7gCaN++fQCAe+65J2b7gQMH8NRTTwEA3nzzTZjNZlRUVCAUCqGsrAx79+7V+2ZkZKChoQGbNm2CqqqYMWMGKisr8fLLL1/fkVB6YgClNZMw4K+YQCAAm80GTdOgKIrscmgKeTv/G72t/6Wv//4Hyf4TmdP4B8lS2URfo/wsGKU24/1+pDgwgCilGfAEneLAAKLUxgBKawwgSm0MoLTGAKKUJsA3IqYzBhClNBH9UwCZzABMUmqhxGMAUUob+teFmHWrciNM/FLCtMEAopQWHRmJWTdlZMBk4hlQumAAkaGYYAIvwdIHA4iMhWc/aYUBRMbCAEorDCAyFBNMvAJLIwwgMhaeAaUVBhAZi4mT0OmEAUSGYmL4pBUGEBkLL8HSCgOIDIVvQkwvDCAyGBPPgtIIA4iMxcRZoHTCACJj4dlPWmEAkaHw/Ce9MIDIWPg+oLTCACJD4V2w9MIAIoNhAKUT/mk5kmpoaAhDQ0PjtkcikZj1cCSCPk0b90woOzsbFosloTXS1GEAkVSHDh3Czp07x23/j3+7A6W3OPT1+vpPsHf9Toz3XRm7d+/GI488kuAqaaowgEiqYDCI8+fPj9veP7gIP125A5cjhSiw/ILgwD9w/vz5cQNocHBwagqlKRHXHNC+ffuwePFiKIoCRVGgqio+//xzvX1oaAhutxszZ85ETk4OKioq4PP5YvbR09OD8vJyZGdno6CgANu3b8fw8HBijobSzv9euR0/Df4eQN0DpfhlsHjc8CHjiSuAZs+ejddeew3t7e04deoU7r33Xjz00EPo6uoCAGzbtg2fffYZ6urq0NLSggsXLsScDo+MjKC8vBzhcBitra04dOgQDh48iB07diT2qCht9A/fgNEfUwEzgiM2uQVRQsUVQA8++CBWr16N+fPn4+abb8arr76KnJwcnDhxApqm4b333sMbb7yBe++9FyUlJThw4ABaW1tx4sQJAMBXX32FM2fO4IMPPsBtt92GVatWYefOndizZw/C4fCUHCAZW6H1J2SaQgAELKZBFGT9U3ZJlECTngMaGRlBXV0dBgYGoKoq2tvbEYlE4HK59D4LFiyA0+mEx+PBihUr4PF4sGjRItjtdr1PWVkZNm3ahK6uLtx+++1x1XD27Fnk5ORM9hAoBfj9/mu2f9/xFXJ7foE2fCPysy7i4sXua/a/cOECzpw5k8gSaRL6+/sn1C/uAOrs7ISqqhgaGkJOTg7q6+tRXFyMjo4OWCwW5OXlxfS32+3wer0AAK/XGxM+o+2jbeMJhUIIhUL6eiAQAABomsb5I4O7cuXKNduPdfwM4OcJ729wcBB9fX3XUxIlwMDAwIT6xR1At9xyCzo6OqBpGj7++GNUVlaipaUl7gLjUVtbi5deeumq7aWlpVAUZUqfm6ZWa2trQvc3b9483HnnnQndJ8Vv9CThr8T9TmiLxYJ58+ahpKQEtbW1WLJkCd5++204HA6Ew+Grfvv4fD44HL+/j8PhcFx1V2x0fbTPWGpqaqBpmr709vbGWzYRpaDr/ihGNBpFKBRCSUkJsrKy0NTUpLd1d3ejp6cHqqoCAFRVRWdnZ8x1f2NjIxRFQXFx8bjPYbVa9Vv/owsRGV9cl2A1NTVYtWoVnE4ngsEgDh8+jGPHjuHLL7+EzWbD+vXrUV1djfz8fCiKgs2bN0NVVaxYsQIAsHLlShQXF+OJJ57Arl274PV68fzzz8PtdsNqtU7JARJR6oorgPx+P5588klcvHgRNpsNixcvxpdffon7778fAPDmm2/CbDajoqICoVAIZWVl2Lt3r/74jIwMNDQ0YNOmTVBVFTNmzEBlZSVefvnlxB4VGcbo2W2i8HNgxmISQhjujaWBQAA2mw2apvFyzOACgQA0TUvY/vLz8zFjxoyE7Y8mZ6KvUX4WjKTinN7fG/8eEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpImU3YBkyGEAAAEAgHJlRDRWEZfm6Ov1fEYMoB+++03AEBRUZHkSojoWoLBIGw227jthgyg/Px8AEBPT881D45iBQIBFBUVobe3F4qiyC7HEDhmkyOEQDAYRGFh4TX7GTKAzObfp65sNht/KCZBURSOW5w4ZvGbyMkBJ6GJSBoGEBFJY8gAslqtePHFF2G1WmWXYigct/hxzKaWSfzVfTIioiliyDMgIkoPDCAikoYBRETSMICISBpDBtCePXswZ84cTJs2DaWlpWhra5NdkjS1tbVYtmwZcnNzUVBQgDVr1qC7uzumz9DQENxuN2bOnImcnBxUVFTA5/PF9Onp6UF5eTmys7NRUFCA7du3Y3h4OJmHIs1rr70Gk8mErVu36ts4ZkkiDObIkSPCYrGI999/X3R1dYkNGzaIvLw84fP5ZJcmRVlZmThw4IA4ffq06OjoEKtXrxZOp1P09/frfZ5++mlRVFQkmpqaxKlTp8SKFSvEnXfeqbcPDw+LhQsXCpfLJb7//ntx9OhRMWvWLFFTUyPjkJKqra1NzJkzRyxevFhs2bJF384xSw7DBdDy5cuF2+3W10dGRkRhYaGora2VWFXq8Pv9AoBoaWkRQgjR19cnsrKyRF1dnd7nhx9+EACEx+MRQghx9OhRYTabhdfr1fvs27dPKIoiQqFQcg8giYLBoJg/f75obGwUd999tx5AHLPkMdQlWDgcRnt7O1wul77NbDbD5XLB4/FIrCx1aJoG4P8/sNve3o5IJBIzZgsWLIDT6dTHzOPxYNGiRbDb7XqfsrIyBAIBdHV1JbH65HK73SgvL48ZG4BjlkyG+jDqpUuXMDIyEvOfDgB2ux1nz56VVFXqiEaj2Lp1K+666y4sXLgQAOD1emGxWJCXlxfT1263w+v16n3GGtPRtnR05MgRfPfdd/j222+vauOYJY+hAoiuze124/Tp0/jmm29kl5LSent7sWXLFjQ2NmLatGmyy/lbM9Ql2KxZs5CRkXHV3QifzweHwyGpqtRQVVWFhoYGfP3115g9e7a+3eFwIBwOo6+vL6b/H8fM4XCMOaajbemmvb0dfr8fd9xxBzIzM5GZmYmWlhbs3r0bmZmZsNvtHLMkMVQAWSwWlJSUoKmpSd8WjUbR1NQEVVUlViaPEAJVVVWor69Hc3Mz5s6dG9NeUlKCrKysmDHr7u5GT0+PPmaqqqKzsxN+v1/v09jYCEVRUFxcnJwDSaL77rsPnZ2d6Ojo0JelS5di7dq1+r85ZkkiexY8XkeOHBFWq1UcPHhQnDlzRmzcuFHk5eXF3I34O9m0aZOw2Wzi2LFj4uLFi/oyODio93n66aeF0+kUzc3N4tSpU0JVVaGqqt4+ekt55cqVoqOjQ3zxxRfixhtv/FvdUv7jXTAhOGbJYrgAEkKId955RzidTmGxWMTy5cvFiRMnZJckDYAxlwMHDuh9rly5Ip555hlxww03iOzsbPHwww+Lixcvxuzn559/FqtWrRLTp08Xs2bNEs8++6yIRCJJPhp5/hxAHLPk4J/jICJpDDUHRETphQFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJ83897GgAfmpyqQAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "200.0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "play(True)[-1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt39]",
   "language": "python",
   "name": "conda-env-pt39-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
